from multiprocessing import Pool, cpu_count, shared_memory
import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
from tqdm.contrib.concurrent import process_map  # 使用process_map来替代imap
import os


# 辅助函数：评估每个组合
def evaluate_combination(params):
    x_num, y_num, z_num = params['x_num'], params['y_num'], params['z_num']
    points_shm_name = params['points_shm_name']
    fault_nodes_shm_name = params['fault_nodes_shm_name']

    # Attach to the existing shared memory blocks for points and fault nodes
    points_shm = shared_memory.SharedMemory(name=points_shm_name)
    fault_nodes_shm = shared_memory.SharedMemory(name=fault_nodes_shm_name)

    # Create arrays from shared memory
    points = np.ndarray(params['points_shape'], dtype=np.float64, buffer=points_shm.buf)
    fault_nodes = np.ndarray(params['fault_nodes_shape'], dtype=np.float64, buffer=fault_nodes_shm.buf)

    # Reconstruct FaultDistanceCalculator in each process
    distance_calculator = FaultDistanceCalculator(fault_nodes)

    grid_nodes = generate_grid_nodes(points, x_num, y_num, z_num)
    average_distance = distance_calculator.calculate_distances(grid_nodes, batch_size=1000)  # 移除tqdm

    # Clean up shared memory
    points_shm.close()
    fault_nodes_shm.close()

    return (x_num, y_num, z_num, average_distance)



class FaultDistanceCalculator:
    def __init__(self, fault_nodes):
        self.tree = cKDTree(fault_nodes)

    def calculate_distances(self, grid_nodes, batch_size=1000):
        distances = []
        num_batches = (len(grid_nodes) + batch_size - 1) // batch_size

        for i in range(0, len(grid_nodes), batch_size):
            batch = grid_nodes[i:i + batch_size]
            batch_distances, _ = self.tree.query(batch, k=1)
            distances.extend(batch_distances)

        return np.mean(distances)
    



def load_points_from_excel(file_path):
    df = pd.read_excel(file_path, header=0)
    if df.shape[1] < 3:
        raise ValueError("Excel文件中必须包含至少三列数据，分别表示X, Y, Z坐标")

    geo_points = df.iloc[:, :3].values
    fault_points = df[df.iloc[:, 0].apply(lambda x: not float(x).is_integer())].iloc[:, :3].values
    print(len(geo_points))
    print(len(fault_points))
    return geo_points, fault_points


def generate_grid_nodes(points, x_num, y_num, z_num):
    x_min, y_min, z_min = points[:, 0].min(), points[:, 1].min(), points[:, 2].min()
    x_max, y_max, z_max = points[:, 0].max(), points[:, 1].max(), points[:, 2].max()

    x_step = (x_max - x_min) / (x_num - 1) if x_num > 1 else 0
    y_step = (y_max - y_min) / (y_num - 1) if y_num > 1 else 0
    z_step = (z_max - z_min) / (z_num - 1) if z_num > 1 else 0

    i, j, k = np.indices((x_num, y_num, z_num))
    grid_nodes = np.array([
        x_min + i * x_step,
        y_min + j * y_step,
        z_min + k * z_step
    ]).reshape(3, -1).T
    return grid_nodes


def optimize_grid(points, fault_nodes, target_cell_size=(10, 15)):
    x_min, y_min, z_min = points[:, 0].min(), points[:, 1].min(), points[:, 2].min()
    x_max, y_max, z_max = points[:, 0].max(), points[:, 1].max(), points[:, 2].max()
    width, height, depth = x_max - x_min, y_max - y_min, z_max - z_min

    x_num_lower = int(np.ceil(width / target_cell_size[1]))
    x_num_upper = int(np.floor(width / target_cell_size[0]))
    y_num_lower = int(np.ceil(height / target_cell_size[1]))
    y_num_upper = int(np.floor(height / target_cell_size[0]))
    z_num_lower = int(np.ceil(depth / target_cell_size[1]))
    z_num_upper = int(np.floor(depth / target_cell_size[0]))

    if x_num_lower > x_num_upper or y_num_lower > y_num_upper or z_num_lower > z_num_upper:
        raise ValueError("给定的目标单元尺寸无法满足网格划分要求，请检查输入数据或调整目标单元尺寸")

    step = 15
    x_num_range = range(max(x_num_lower, 2), x_num_upper + 1,step)
    y_num_range = range(max(y_num_lower, 2), y_num_upper + 1,step)
    z_num_range = range(max(z_num_lower, 2), z_num_upper + 1,step)
    print(f"x_num_range:{x_num_range},y_num_range:{y_num_range},z_num_range:{z_num_range}")
    combinations = [(x_num, y_num, z_num) for x_num in x_num_range 
                    for y_num in y_num_range 
                    for z_num in z_num_range]

    # 使用共享内存存储点数据和故障节点
    points_shm = shared_memory.SharedMemory(create=True, size=points.nbytes)
    points_shared = np.ndarray(points.shape, dtype=np.float64, buffer=points_shm.buf)
    points_shared[:] = points[:]

    fault_nodes_shm = shared_memory.SharedMemory(create=True, size=fault_nodes.nbytes)
    fault_nodes_shared = np.ndarray(fault_nodes.shape, dtype=np.float64, buffer=fault_nodes_shm.buf)
    fault_nodes_shared[:] = fault_nodes[:]

    params_list = [{
        'x_num': comb[0], 'y_num': comb[1], 'z_num': comb[2],
        'points_shm_name': points_shm.name, 'points_shape': points.shape,
        'fault_nodes_shm_name': fault_nodes_shm.name, 'fault_nodes_shape': fault_nodes.shape
    } for comb in combinations]

    # 使用process_map并指定一个主进度条
    results = list(process_map(evaluate_combination, params_list, max_workers=cpu_count(), total=len(combinations), desc="Evaluating Combinations"))

    # Clean up shared memory
    points_shm.close()
    points_shm.unlink()
    fault_nodes_shm.close()
    fault_nodes_shm.unlink()

    best_x_num, best_y_num, best_z_num, min_average_distance = min(results, key=lambda r: r[3])

    return best_x_num, best_y_num, best_z_num, min_average_distance

def construct_hexahedral_grid(x_num, y_num, z_num):
    if x_num < 2 or y_num < 2 or z_num < 2:
        raise ValueError("The grid dimensions must be at least 2 in each direction.")

    cells = []

    for i in range(x_num - 1):
        for j in range(y_num - 1):
            for k in range(z_num - 1):
                # Using 3D indexing to avoid complex linear indexing
                p0 = (i, j, k)
                p1 = (i + 1, j, k)
                p2 = (i + 1, j + 1, k)
                p3 = (i, j + 1, k)
                p4 = (i, j, k + 1)
                p5 = (i + 1, j, k + 1)
                p6 = (i + 1, j + 1, k + 1)
                p7 = (i, j + 1, k + 1)

                # Convert 3D indices back to linear indices
                def to_linear(ijk):
                    return ijk[0] * y_num * z_num + ijk[1] * z_num + ijk[2]

                cell = [to_linear(p) for p in [p0, p1, p2, p3, p4, p5, p6, p7]]

                # Check for duplicate indices in the cell before appending
                if len(set(cell)) != 8:
                    print(f"Duplicate indices found at i={i}, j={j}, k={k}: {cell}")
                    raise ValueError(f"Duplicate point indices found in cell: {cell}")

                cells.append(cell)

    return np.array(cells)

def export_to_vtk(points, cells, fault_points, output_file):
    num_points = len(points)
    num_cells = len(cells)
    num_fault_points = len(fault_points)

    with open(output_file, "w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("Hexahedral grid with fault surface\n")
        f.write("ASCII\n")
        f.write("DATASET UNSTRUCTURED_GRID\n")
        
        # Write POINTS
        f.write(f"POINTS {num_points + num_fault_points} float\n")
        for point in points:
            f.write(f"{point[0]} {point[1]} {point[2]}\n")
        for fault_point in fault_points:
            f.write(f"{fault_point[0]} {fault_point[1]} {fault_point[2]}\n")

        # Write CELLS
        f.write(f"CELLS {num_cells + num_fault_points} {num_cells * 9 + num_fault_points * 2}\n")
        for cell in cells:
            f.write(f"8 {' '.join(map(str, cell))}\n")
        for i in range(num_fault_points):  # Writing fault points as individual vertices
            f.write(f"1 {i + num_points}\n")  # Assuming fault points are added after grid points

        # Write CELL_TYPES (12 indicates a hexahedron, 1 indicates a vertex in VTK)
        f.write(f"CELL_TYPES {num_cells + num_fault_points}\n")
        for _ in range(num_cells):
            f.write("12\n")
        for _ in range(num_fault_points):
            f.write("1\n")

        # Optionally, you can add CELL_DATA or POINT_DATA if you have scalar/vector data to visualize

    print(f"VTK 文件已成功导出到 {output_file}")

# 修改主程序调用，加入故障点参数
if __name__ == "__main__":
    input_file = r"D:\360MoveData\Users\10758\Desktop\first_code\000035\q3_result\GridE_change.xlsx"
    output_file = r"D:\360MoveData\Users\10758\Desktop\first_code\000035\q3_result\output_lgr.vtk"
    
    print("加载数据")
    geo_points, fault_points = load_points_from_excel(input_file)
    print("网格生成参数")
    target_cell_size = (10, 15)
    best_x_num, best_y_num, best_z_num, min_average_distance = optimize_grid(geo_points, fault_points, target_cell_size)
    print(f"最佳网格参数: x_num={best_x_num}, y_num={best_y_num}, z_num={best_z_num}, 最小平均距离: {min_average_distance}")

    # 生成网格节点
    grid_nodes = generate_grid_nodes(geo_points, best_x_num, best_y_num, best_z_num)

    # 使用最优参数重新构建六面体单元
    cells = construct_hexahedral_grid(best_x_num, best_y_num, best_z_num)

    # 导出为 VTK 文件，现在包括了断层信息
    export_to_vtk(grid_nodes, cells, fault_points, output_file)